import os
from pprint import pprint

from dotenv import load_dotenv
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import create_sql_agent, SQLDatabaseToolkit
from langchain_openai import AzureChatOpenAI

"""
代码说明：
首先创建/连接到SQLite数据库并设置示例数据
通过Ollama初始化本地Llama3.2模型 
 (1) 报错：ValueError: This function requires a bind_tools() method be implemented on the LLM.意思是：因为create_sql_agent需要LLM支持工具调用功能引起的。
（2）解决一：换个模型如下。
（3）解决二：自定义一个sql 工具
使用create_sql_agent创建能够处理SQL查询的代理
代理可以理解自然语言问题并转换为SQL查询
执行查询并返回自然语言格式的结果
"""

# 1. 初始化SQLite数据库连接
# 连接到SQLite数据库文件（如果不存在会自动创建）
db = SQLDatabase.from_uri("sqlite:///example.db")


# 2. 创建示例数据表并插入一些数据
# 在实际应用中你可能已经有现有的数据库
def setup_database():
    import sqlite3
    conn = sqlite3.connect('example.db')
    cursor = conn.cursor()

    # 创建用户表
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS users (
        id INTEGER PRIMARY KEY,
        name TEXT NOT NULL,
        age INTEGER,
        email TEXT
    )
    ''')

    # 插入示例数据
    cursor.execute("INSERT OR IGNORE INTO users (name, age, email) VALUES ('Alice', 30, 'alice@email.com')")
    cursor.execute("INSERT OR IGNORE INTO users (name, age, email) VALUES ('Bob', 25, 'bob@email.com')")
    cursor.execute("INSERT OR IGNORE INTO users (name, age, email) VALUES ('Charlie', 35, 'charlie@email.com')")

    conn.commit()
    conn.close()


# 初始化数据库
setup_database()

# 3. 初始化Llama3.2模型（通过Ollama）
# 确保你已经在本地运行了Ollama并拉取了llama3.2:3b模型
# llm = OllamaLLM(model="llama3.2:3b")

load_dotenv()

llm = AzureChatOpenAI(
    # openai_api_key=
    # openai_api_base=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
    api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
    temperature=0.7
)

toolkit = SQLDatabaseToolkit(db=db,llm=llm)
print("\n查询工具结果:")
pprint(toolkit.get_tools())

# 4. 创建SQL代理
agent = create_sql_agent(
    llm=llm,
    # db=db,
    toolkit=toolkit,
    agent_type="tool-calling",  # 使用工具调用代理
    verbose=False  # 显示详细执行过程
)

# 5. 运行查询
try:
    # 示例查询：询问用户信息
    result = agent.invoke("查询所有用户的姓名和年龄")
    print("\n查询结果:")
    print(result["output"])

    # 另一个示例：条件查询
    result2 = agent.invoke("年龄大于28岁的用户有哪些？")
    print("\n条件查询结果:")
    print(result2["output"])

except Exception as e:
    print(f"执行出错: {e}")
